import numpy as np
from Kohonen.fuzzydist import fuzzyDist


def PSO(x, w1, K):
    # 参数初始化
    # 粒子群算法中的两个参数
    c1 = 1.49445
    c2 = 1.49445

    maxGen = 2  # 进化次数
    sizePop = K  # 种群规模

    v_max = 0.1
    V_min = -0.1
    popMax = 1
    popMin = 0
    row, col = w1.shape
    # 产生初始粒子和速度
    pop = np.zeros([sizePop, row])
    V = np.zeros([sizePop, row])
    fitness = np.zeros(sizePop)
    for i in range(sizePop):
        # 随机产生一个种群
        pop[i, :] = np.transpose(w1[:, i])    # 初始种群
        V[i, :] = np.random.rand(1, row)  # 初始化速度
        # 计算适应度
        fitness[i] = fuzzyDist(x, pop[i])  # 染色体的适应度
    # 个体极值和群体极值
    bestFitness = np.min(fitness)
    bestIndex = np.where(fitness == bestFitness)
    zBest = pop[bestIndex, :]  # 全局最佳
    gBest = pop  # 个体最佳
    fitnessGBest = fitness  # 个体最佳适应度值
    fitnessZBest = bestFitness  # 全局最佳适应度值
    # 迭代寻优
    for i in range(maxGen):
        for j in range(sizePop):
            # 速度更新
            V[j, :] = V[j, :] + c1 * np.random.rand() * (gBest[j, :] - pop[j, :]) + c2 * np.random.rand() * (zBest - pop[j, :])
            V[j, np.where(V[j, :] > v_max)] = v_max
            V[j, np.where(V[j, :] < V_min)] = V_min
            # 种群更新
            pop[j, :] = pop[j, :] + 0.5 * V[j, :]
            pop[j, np.where(pop[j, :] > popMax)] = popMax
            pop[j, np.where(pop[j, :] < popMin)] = popMin
            # 适应度值
            fitness[j] = fuzzyDist(x, pop[j, :])

        for j in range(sizePop):
            # 个体最优更新
            if fitness[j] < fitnessGBest[j]:
                gBest[j, :] = pop[j, :]
                fitnessGBest[j] = fitness[j]
            # 群体最优更新
            if fitness[j] < fitnessZBest:
                zBest = pop[j, :]
                fitnessZBest = fitness[j]

    return fitnessGBest, np.transpose(pop)